import segyio
import json

from utils.index import resource_path


# 从指定的起始索引开始，查找第一个非零值的时间索引
def find_first_non_zero_time(data, start_index):
    """
    从指定的起始索引开始，查找第一个非零值的时间索引。
    :param data: 数据数组
    :param start_index: 起始索引
    :return: 第一个非零值的时间索引，如果没有非零值，返回None
    """
    for i in range(start_index, len(data)):
        if data[i] != 0:
            return i
    return None


# 判断sgy文件的数据是向左偏移还是向右偏移
def get_flag(segy_file, start_trace, start_time_idx=375):
    """
        判断sgy文件的数据是向左偏移还是向右偏移
        :param segy_file: SEG-Y文件数据
        :param start_trace: 起始道编号
        :param start_time_idx: 第一道开始的时间
        :return: 偏移标志，如果是-1，则向左偏移，如果是1，则向右偏移
        """
    flag = -1  # 如果是-1，则向左偏移，如果是1，则向右偏移

    # 从指定道开始处理
    trace_index = start_trace - 1  # segyio的道索引从0开始
    trace_data = segy_file.trace[trace_index]  # 读取当前道的数据
    start_time_index = start_time_idx # int(1500 / 4) # 第1道开始时间为1500，避免数据的问题
    non_zero_time_index = find_first_non_zero_time(trace_data, start_time_index)  # 查找第一个非零值的时间索引
    # 左侧第2道数据
    trace_index2 = trace_index + flag
    trace_data2 = segy_file.trace[trace_index2]  # 读取当前道的数据
    start_time_index2 = non_zero_time_index  # 开始时间为上一道第一个不为0的时间索引
    non_zero_time_index2 = find_first_non_zero_time(trace_data2, start_time_index2)  # 查找第一个非零值的时间索引
    # 判断左偏还是右偏
    if non_zero_time_index == non_zero_time_index2:
        flag = 1

    return flag


# 处理SEG-Y文件，从指定道开始，逐道查找第一个非零值的时间
def get_trace_offset(segy_file, start_trace, trace_num, start_time_idx=370):
    """
    处理SEG-Y文件，从指定道开始，逐道查找第一个非零值的时间。
    :param segy_file: SEG-Y文件数据
    :param start_trace: 起始道编号
    :param trace_num: 走廊通道数
    :return: 道偏移的相关数据，
    [
        {
        "trace_index"，道编号，从1开始
        "non_zero_time"，该道第一个非零点的时间
        },
        ...
    ]
    """
    sample_interval = segy_file.bin[segyio.BinField.Interval] / 1000  # 获取采样间隔（ms），转换为ms
    flag = get_flag(segy_file, start_trace, start_time_idx)  # 获取偏移方向
    trace_offset = []  # 存储道偏移的相关数据

    if flag == 1:  # 向右偏移时，开始计算的初始道为最右侧的道
        start_trace = start_trace + trace_num + 1
        print("=====*****=====\n向右偏移")
    else:
        print("=====*****=====\n向左偏移")

    print(f"=====*****=====\n偏移起始道：{start_trace}")

    # 从指定道开始处理
    # trace_index = start_trace + flag  # segyio的道索引从0开始
    trace_index = start_trace - 1  # segyio的道索引从0开始，所以要-1
    # trace_index = 160
    last_non_zero_time = 0

    while trace_index >= 0:
        trace_data = segy_file.trace[trace_index]  # 读取当前道的数据
        if trace_index == start_trace - 1:  # 如果是第一道，从0ms开始排查
            start_time_index = start_time_idx  # int(1500 / 4)
        else:  # 否则从上一道记录的时间开始排查
            start_time_index = last_non_zero_time

        # 查找第一个非零值的时间索引
        non_zero_time_index = find_first_non_zero_time(trace_data, start_time_index)
        if non_zero_time_index is not None:  # 有非零值
            non_zero_time = non_zero_time_index * sample_interval  # 通过索引和采样间隔得到事件
            print(f"道 {trace_index + 1}: 第一个非零值的时间为 {non_zero_time} ms")
            last_non_zero_time = non_zero_time_index
            trace_offset.append({
                "trace_index": trace_index + 1,  # 道编号
                "non_zero_time": non_zero_time  # 该道第一个非零点的时间
            })
        else:  # 没有找到非零值
            print(f"道 {trace_index + 1}: 没有找到非零值")
            break

        # 准备处理下一道
        trace_index = trace_index + flag

    return trace_offset


# 将走廊中的道数据扩充到指定道数
def create_trace_by_num(inp_path, out_file, trace_num):
    """
    将走廊中的道数据扩充到指定道数
    :param inp_path: 为走廊数据的文件地址
    :param out_path: 扩充后的走廊数据保存的文件地址
    :param trace_num: 指定扩充道数
    """
    with segyio.open(inp_path, 'r', ignore_geometry=True) as infile:
        # 以写入模式打开输出的 SGY 文件
        spec = segyio.spec()
        spec.sorting = infile.sorting
        spec.format = infile.format
        spec.samples = infile.samples
        spec.tracecount = trace_num  # 数据总道数
        # out_file = r"E:\vspData\斜井\repeat_st.sgy"
        with segyio.create(out_file, spec) as outfile:
            # 复制二进制头信息
            outfile.bin = infile.bin
            # 复制道头信息和数据
            for i in range(trace_num):
                outfile.header[i] = infile.header[0]
                outfile.trace[i] = infile.trace[0]


# 通过裁剪得到地面数据，包括start_index和end_index
def clip_trace(inp_file, out_file, start_index, end_index):
    """
    通过裁剪得到走廊数据，包括start_index和end_index
    :param inp_file: 被截取的sgy文件路径，地面镶嵌数据
    :param out_file: 截取后保存的sgy文件路径，走廊
    :param start_index: 截取道的起始下标，地面镶嵌数据中走廊的起始位置
    :param end_index: 截取道的结束下标，地面镶嵌数据中走廊的结束位置
    """
    with segyio.open(inp_file, 'r', ignore_geometry=True) as infile:
        # 以写入模式打开输出的 SGY 文件
        spec = segyio.spec()
        spec.sorting = infile.sorting
        spec.format = infile.format
        spec.samples = infile.samples
        spec.tracecount = end_index - start_index + 1  # 数据总道数
        with segyio.create(out_file, spec) as outfile:
            # 复制二进制头信息
            outfile.bin = infile.bin
            # 复制道头信息和数据
            for i, trace_index in enumerate(range(start_index - 1, end_index)):
                outfile.header[i] = infile.header[trace_index]
                outfile.trace[i] = infile.trace[trace_index]


# 获取偏移数据，生成指定道数的走廊数据和滤波走廊数据，截取地面数据
def get_inlay_data(file_path, sei_path, sei_num, stk_path, stk_path2, stkbf_path, stkbf_path2, trace_start, trace_num=10):
    """
    获取偏移数据，生成指定道数的走廊数据和滤波走廊数据，截取地面数据
    :param file_path: 镶嵌数据，含走廊和滤波走廊
    :param sei_path: 地面数据，保存路径
    :param sei_num: 地面数据需要截取的道数
    :param stk_path: 走廊数据
    :param stk_path2: 指定道数的走廊数据，保存路径
    :param stkbf_path: 滤波走廊数据
    :param stkbf_path2: 指定道数的滤波走廊数据，保存路径
    :param trace_start: 滤波走廊开始道号
    :param trace_num: 走廊总道数，原始
    :return:
    """
    with open(resource_path("static/json/setting.json"), 'r', encoding='utf-8') as file:  # 读取设置文件
        setting = json.load(file)
        first_start_time = setting["first_start_time"]
        print(f"=====*****=====\n起始道计算非零点的起始时间：{first_start_time}ms")
    with segyio.open(file_path, strict=False, ignore_geometry=True) as segy_file:
        dt = segy_file.bin[segyio.BinField.Interval] / 1000  # 采样间隔
        print(f"=====*****=====\n采样间隔：{dt}")
        trace_offset = get_trace_offset(segy_file, trace_start, trace_num, int(first_start_time / dt))
        flag = get_flag(segy_file, trace_start, int(first_start_time / dt))

        trace_num_offset = len(trace_offset) + trace_num  # 新的走廊总道数
        # 创建新的走廊数据
        create_trace_by_num(stk_path, stk_path2, trace_num_offset)
        # 创建新的滤波走廊数据
        create_trace_by_num(stkbf_path, stkbf_path2, trace_num_offset)

        # 截取地面
        if flag == -1:  # 向左偏移
            start_index = trace_offset[len(trace_offset) - 1]["trace_index"]  # 截取道的开始下标，392
            end_index = trace_start + trace_num  # 截取道的结束下标，414
        else:  # 向右偏移
            start_index = trace_start + 1  # 截取道的开始下标，405
            end_index = trace_offset[len(trace_offset) - 1]["trace_index"]  # 截取道的结束下标,427
        # clip_trace(file_path, stkbf_path, start_index, end_index)  # 裁剪得到滤波走廊数据
        start_index2 = start_index + trace_num + 1  # 截取道的开始下标，403，416
        end_index2 = end_index + trace_num + 1  # 截取道的结束下标，425，438
        # clip_trace(file_path, stk_path, start_index2, end_index2)  # 裁剪得到滤波走廊数据
        start_index3 = start_index2 + trace_num + 1  # 截取道的开始下标，414，427
        end_index3 = end_index2 + sei_num + 1  # 截取道的结束下标，476，489
        clip_trace(file_path, sei_path, start_index3, end_index3)  # 裁剪得到滤波走廊数据

    # 偏移后开始道号
    if flag == -1:  # 向左偏移
        start_trace2 = len(trace_offset)  # 起始道，13
    else:
        start_trace2 = 0  # 起始道，0

    # 截取道的开始下标
    if flag == -1:  # 向左偏移
        start_index = trace_offset[len(trace_offset) - 1]["trace_index"]  # 截取道的开始下标，392
    else:  # 向右偏移
        start_index = trace_start + 1  # 截取道的开始下标，405

    # 将地面嵌入数据的相关偏移数据，改为裁剪出来的走廊的相关偏移数据
    for i in range(len(trace_offset)):
        trace_offset[i]["trace_index"] = trace_offset[i]["trace_index"] - start_index + 1

    return {
        "trace_start_offset": start_trace2,  # 走廊起始道编号，偏移后的，滤波走廊和走廊相同，左偏为13，右偏为0
        "trace_num_offset": trace_num_offset,  # 总道数，偏移后的，滤波走廊和走廊相同
        "flag": flag,  # 标志，左偏还是右偏
        "trace_offset": trace_offset,  # 偏移相关数据，偏移后的，滤波走廊和走廊相同，左侧第一道为13，右侧第一道为11
        "trace_num": trace_num  # 总道数，偏移前的
    }


if __name__ == "__main__":
    file_path = r"E:\vspData\斜井\psdm_stk_inl994_1394_0914vsp_mu_filter_scale_inq_3D_tm_inlne1194AGC镶嵌1.sgy"
    stk_path = r"E:\vspData\斜井\Tk1_最终拉平1ms_COR100.sgy"
    stk_path2 = r"E:\vspData\斜井\repeat_stk.sgy"
    stkbf_path = r"E:\vspData\斜井\Tk1_最终拉平1ms_COR100LB.sgy"
    stkbf_path2 = r"E:\vspData\斜井\repeat_stkbp.sgy"
    trace_num = 10
    trace_start = 404

    # with segyio.open(file_path, strict=False, ignore_geometry=True) as segy_file:
    #     trace_offset = get_trace_offset(segy_file, trace_start, trace_num)
    #     flag = get_flag(segy_file, trace_start)
    #
    #     trace_start2 = len(trace_offset) + trace_num  # 新的走廊总道数
    #     create_trace_by_num(stk_path, stk_path2, len(trace_offset) + trace_num)  # 创建新的走廊数据
    #     create_trace_by_num(stkbf_path, stkbf_path2, len(trace_offset) + trace_num)  # 创建新的走廊数据
